# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/24 17:54
# @Author  : Dell
# @File    : back_question.py
# @Software: PyCharm
# @Desc    :问题回退，对太多/太过细节/太繁琐的问题做下总结
from langchain.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(model_name='qwen-plus', base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                     temperature=0,
                     api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37")
# ----------------------- 配置阶段 -----------------------
examples = [
    {
        "input": "警察进行合法逮捕吗？",
        "output": "警察可以做什么？",
    },
    {
        "input": "Jan Sindel’s was born in what country?",
        "output": "what is Jan Sindel’s personal history?",
    },
]

# 正确：使用已实现抽象方法的 ChatPromptTemplate
example_prompt = ChatPromptTemplate.from_messages([
    ("human", "{input}"),  # 明确消息角色
    ("ai", "{output}")
])
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples,
)
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """您是世界知识方面的专家。您的任务是退一步思考，将问题解释为更通用的退一步思考问题，这样更容易回答。以下是几个示例：""",
        ),
        # 一些镜头示例
        few_shot_prompt,
        # 新的问题
        ("user", "{question}"),
    ]
)
# ----------------------- 执行阶段 -----------------------
generate_queries_step_back = prompt | llm | StrOutputParser()
question = "人工智能如何用于金融风控？"
back_re=generate_queries_step_back.invoke({"question": question})
print(f"打印回退问题：{back_re}")
# 响应提示
response_prompt_template = """您是世界知识专家。我要问您一个问题。您的回答应全面，并且不与以下上下文相矛盾（如果它们相关）。否则，如果它们不相关，请忽略它们。

# {normal_context}
# {step_back_context}

# 原始问题：{question}
# 答案："""
response_prompt = ChatPromptTemplate.from_template(response_prompt_template)

chain = (
    {
        # 使用普通问题检索上下文
        "normal_context": RunnableLambda(lambda x: x["question"]) ,
        # 使用后退问题检索上下文
        "step_back_context": generate_queries_step_back,
        # 把问题转过去
        "question": lambda x: x["question"],
    }
    | response_prompt
    | ChatOpenAI(temperature=0)
    | StrOutputParser()
)

f_re=chain.invoke({"question": question})
print(f"打印回退：{f_re}")